package cn.tworice;

import cn.tworice.source.JavaClassObject;

import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileManager;
import javax.tools.JavaFileObject;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URI;

import javax.tools.*;
import java.util.HashMap;
import java.util.Map;

public class MemoryJavaFileManager extends ForwardingJavaFileManager<JavaFileManager> {
    // ⭐ 使用Map存储编译生成的字节码
    private final Map<String, byte[]> classBytes = new HashMap<>();

    public MemoryJavaFileManager(JavaFileManager fileManager) {
        super(fileManager);
    }

    @Override
    public JavaFileObject getJavaFileForOutput(
            Location location,
            String className,
            JavaFileObject.Kind kind,
            FileObject sibling) throws IOException {
        // ⭐ 生成内存中的JavaFileObject
        if (kind == JavaFileObject.Kind.CLASS) {
            return new MemoryJavaFileObject(className);
        }
        return super.getJavaFileForOutput(location, className, kind, sibling);
    }

    // ⭐ 新增：获取存储的字节码
    public Map<String, byte[]> getClassBytes() {
        return new HashMap<>(classBytes);
    }

    // ⭐ 自定义JavaFileObject实现
    private class MemoryJavaFileObject extends SimpleJavaFileObject {
        private final String className;

        MemoryJavaFileObject(String className) {
            super(URI.create("bytes:///" + className.replace('.', '/') + Kind.CLASS.extension), Kind.CLASS);
            this.className = className;
        }

        @Override
        public OutputStream openOutputStream() {
            return new ByteArrayOutputStream() {
                @Override
                public void close() throws IOException {
                    // ⭐ 将字节码存入Map
                    classBytes.put(className, this.toByteArray());
                }
            };
        }
    }
}